from flask import Flask, render_template, request, redirect, url_for, flash
from flask_sqlalchemy import SQLAlchemy
from flask_login import LoginManager, UserMixin, login_user, login_required, logout_user, current_user
from werkzeug.security import generate_password_hash, check_password_hash
import os

app = Flask(__name__)
app.config['SECRET_KEY'] = 'your-secret-key'
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///chat.db'
db = SQLAlchemy(app)
login_manager = LoginManager()
login_manager.init_app(app)
login_manager.login_view = 'login'

class User(UserMixin, db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True, nullable=False)
    password = db.Column(db.String(120), nullable=False)
    role = db.Column(db.String(20), nullable=False)  # 'teacher' or 'student'
    messages = db.relationship('Message', backref='author', lazy=True)

class Message(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    content = db.Column(db.Text, nullable=False)
    user_id = db.Column(db.Integer, db.ForeignKey('user.id'), nullable=False)
    timestamp = db.Column(db.DateTime, nullable=False, default=db.func.current_timestamp())
    conversation_id = db.Column(db.Integer, db.ForeignKey('conversation.id'), nullable=False)

class Conversation(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    student_id = db.Column(db.Integer, db.ForeignKey('user.id'), nullable=False)
    teacher_id = db.Column(db.Integer, db.ForeignKey('user.id'), nullable=False)
    messages = db.relationship('Message', backref='conversation', lazy=True)

@login_manager.user_loader
def load_user(user_id):
    return User.query.get(int(user_id))

@app.route('/')
def index():
    return render_template('login.html')

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        role = request.form.get('role')
        
        user = User.query.filter_by(username=username).first()
        if user and check_password_hash(user.password, password) and user.role == role:
            login_user(user)
            return redirect(url_for('chat'))
        flash('Invalid username, password or role')
    return render_template('login.html')

@app.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        role = request.form.get('role')
        
        if User.query.filter_by(username=username).first():
            flash('Username already exists')
            return redirect(url_for('register'))
        
        hashed_password = generate_password_hash(password)
        new_user = User(username=username, password=hashed_password, role=role)
        db.session.add(new_user)
        db.session.commit()
        
        return redirect(url_for('login'))
    return render_template('register.html')

@app.route('/chat')
@app.route('/chat/<int:conversation_id>')
@login_required
def chat(conversation_id=None):
    if current_user.role == 'student':
        conversations = Conversation.query.filter_by(student_id=current_user.id).all()
    else:
        conversations = Conversation.query.filter_by(teacher_id=current_user.id).all()
    
    # 如果指定了会话ID，获取该会话
    current_conversation = None
    if conversation_id:
        current_conversation = Conversation.query.get(conversation_id)
    # 如果没有指定会话ID但有会话列表，使用第一个会话
    elif conversations:
        current_conversation = conversations[0]
    
    return render_template('chat.html', 
                         conversations=conversations,
                         conversation=current_conversation)

@app.route('/create_conversation', methods=['POST'])
@login_required
def create_conversation():
    if current_user.role != 'student':
        flash('只有学生可以创建新会话')
        return redirect(url_for('chat'))
    
    # 获取一个随机教师
    teacher = User.query.filter_by(role='teacher').first()
    if not teacher:
        flash('没有可用的教师')
        return redirect(url_for('chat'))
    
    conversation = Conversation(student_id=current_user.id, teacher_id=teacher.id)
    db.session.add(conversation)
    db.session.commit()
    
    return redirect(url_for('chat', conversation_id=conversation.id))

@app.route('/send_message', methods=['POST'])
@login_required
def send_message():
    content = request.form.get('content')
    conversation_id = request.form.get('conversation_id')
    
    if content and conversation_id:
        message = Message(content=content, user_id=current_user.id, conversation_id=conversation_id)
        db.session.add(message)
        db.session.commit()
    
    return redirect(url_for('chat', conversation_id=conversation_id))

@app.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('login'))

if __name__ == '__main__':
    with app.app_context():
        db.create_all()
    app.run(debug=True) 